# 导包
import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_curve, roc_auc_score
# 标准化
from sklearn.preprocessing import StandardScaler
# XGBoost
from xgboost import XGBClassifier
# 交叉网格
from sklearn.model_selection import GridSearchCV

# 读取数据
data = pd.read_csv('../data/train.csv')

# 删除无用列
data = data.drop(['StandardHours'], axis=1)

#  数据预处理
x = data.drop(['Attrition'], axis=1)
y = data['Attrition']
x = pd.get_dummies(x)
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42)
# 标准化
scaler = StandardScaler()
x_train = scaler.fit_transform(x_train)
x_test = scaler.transform(x_test)
model = RandomForestClassifier(bootstrap=False, max_depth=10, min_samples_leaf=1, min_samples_split=10, n_estimators=50)
model.fit(x_train, y_train)
# 模型评估
print(model.score(x_test, y_test))
print('-----------------------------')
print(model.feature_importances_)
# 生成ROC曲线
# 获取预测概率（注意是 predict_proba）
y_pred_proba = model.predict_proba(x_test)[:, 1]

# 计算 FPR, TPR 和阈值
fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)

# 计算 AUC 分数
auc_score = roc_auc_score(y_test, y_pred_proba)

# 绘制 ROC 曲线
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, label=f'ROC Curve (AUC = {auc_score:.2f})')
plt.plot([0, 1], [0, 1], 'k--', label='Random Guess')  # 对角线
plt.xlabel('False Positive Rate (FPR)')
plt.ylabel('True Positive Rate (TPR)')
plt.title('Receiver Operating Characteristic (ROC) Curve - Random Forest')
plt.legend(loc='lower right')
plt.grid()
plt.show()

# 生成权重与对应的特征名称以柱状图的形式表现出来
# 画布再调大一点
plt.figure(figsize=(20, 10))
feature_importances = pd.Series(model.feature_importances_, index=x.columns)
feature_importances.nlargest(20).plot(kind='barh')
plt.show()



# 提取20个特征
x = x[feature_importances.nlargest(20).index]
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=0)
model = XGBClassifier(n_estimators=100, learning_rate=0.1, max_depth=5, min_child_weight=1, gamma=0, subsample=0.8)
model.fit(x_train, y_train)
print(model.score(x_test, y_test))
# 获取预测概率（注意是 predict_proba）
y_pred_proba = model.predict_proba(x_test)[:, 1]
# 计算 FPR, TPR 和阈值
fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
# 计算 AUC 分数
auc_score = roc_auc_score(y_test, y_pred_proba)
# 绘制 ROC 曲线
plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, label=f'ROC Curve (AUC = {auc_score:.2f})')
plt.plot([0, 1], [0, 1], 'k--', label='Random Guess')  # 对角线
plt.xlabel('False Positive Rate (FPR)')
plt.ylabel('True Positive Rate (TPR)')
plt.title('Receiver Operating Characteristic (ROC) Curve - XGBoost')
plt.legend(loc='lower right')
plt.grid()
plt.show()
# 保存模型至models
import os
import joblib

# 确保 models 文件夹存在
os.makedirs('models', exist_ok=True)
# 保存 XGBoost 模型
model_path = os.path.join('models', 'xgboost_model.pkl')
joblib.dump(model, model_path)
print(f"XGBoost 模型已保存至: {model_path}")








